import openpyxl
import pandas as pd
import os

def read_file(excel_file):
    df = pd.read_excel(excel_file)
    print(df)

def read_txt_file(txt_file):
    """从 txt 文件中读取数据，跳过无关内容并从最后一个 'Overall' 开始提取"""
    data = {}
    start_reading = False  # 标志位，指示是否开始读取数据
    last_overall_line = None  # 用于记录最后一个 'Overall' 出现的位置
    
    try:
        with open(txt_file, 'r') as f:
            lines = f.readlines()  # 一次性读取所有行
            
            # 从最后一行向上扫描，寻找最后一个 'Overall'
            for i in range(len(lines) - 1, -1, -1):
                line = lines[i].strip()
                if "Overall" in line:
                    last_overall_line = i
                    break  # 找到最后一个 'Overall' 后停止搜索
        
            if last_overall_line is not None:
                # 从最后一个 'Overall' 之后开始读取数据
                for line in lines[last_overall_line + 1:]:
                    line = line.strip()
                    
                    # 假设文件中的每一行是 'a: num' 的格式
                    if ':' in line:
                        key, value = line.split(':')
                        data[key.strip()] = float(value.strip())  # 转换为浮动数值
            else:
                print("Error: 'Overall' not found in the file.")
                
    except FileNotFoundError:
        print(f"Error: {txt_file} not found.")
    
    return data

def find_empty_row(sheet):
    """寻找 Excel 表格中的第一空白行"""
    for row_idx in range(2, sheet.max_row + 2):  # 从第二行开始
        if not sheet.cell(row=row_idx, column=1).value:  # 如果第一列为空
            return row_idx
    return sheet.max_row + 1  # 如果所有行都填充了数据，返回下一行

def write_to_excel(excel_file, txt_data):
    """将数据写入 Excel，txt_data 对应于 Excel 表格中的列标题"""
    try:
        # 打开现有的 Excel 文件
        workbook = openpyxl.load_workbook(excel_file)
        sheet = workbook.active

        # 获取 Excel 表格的第一行（列标题）
        headers = [cell.value for cell in sheet[1]]  # 获取第一行的所有列名

        print(f"Excel headers: {headers}")

        empty_row = find_empty_row(sheet)
        print(f"Writing to empty row: {empty_row}")

        # 根据 txt_data 中的 key 填写相应的列
        for key, num in txt_data.items():
            if key in headers:
                # print(f"key {key} is searched, value is {num}.\n")
                # 获取列标题在第一行的位置，返回索引从 1 开始
                col_idx = headers.index(key) + 1
                # 按行填写数据
                row = empty_row  # 从第二行开始填写
                print(f"Writing {key}: {num} to column {col_idx}")
                # while row <= sheet.max_row:
                sheet.cell(row=row, column=col_idx, value=num)
                    # row += 1

        # 保存修改后的 Excel 文件
        workbook.save(excel_file)
        print(f"Data has been written to {excel_file}")
        workbook.close()

    except Exception as e:
        print(f"Error while writing to Excel: {e}")

def process_folder(folder_path):
    """处理文件夹中的所有 txt 文件"""
    folder_path1 = folder_path + "/12_3"
    for file_name in os.listdir(folder_path1):
        if file_name.startswith('llama3.1-70b'):
            full_path = os.path.join(folder_path1, file_name)
            txt_data = read_txt_file(full_path)
            read_file(full_path)

            if "SLO" in file_name:
                if "prefix" in file_name:
                    excel_file = os.path.join(folder_path, "result_llama3.1-70b_SLO_prefix.xlsx")
                elif "select" in file_name:
                    excel_file = os.path.join(folder_path, "result_llama3.1-70b_SLO_selectbatch.xlsx")
                elif "splitifuse" in file_name:
                    excel_file = os.path.join(folder_path, "result_llama3.1-70b_SLO_splitfuse.xlsx")
                else:
                    excel_file = os.path.join(folder_path, "result_llama3.1-70b_SLO.xlsx")
            else:
                if "prefix" in file_name:
                    excel_file = os.path.join(folder_path, "result_llama3.1-70b_decode_prefix.xlsx")
                elif "select" in file_name:
                    excel_file = os.path.join(folder_path, "result_llama3.1-70b_decode_selectbatch.xlsx")
                elif "splitifuse" in file_name:
                    excel_file = os.path.join(folder_path, "result_llama3.1-70b_decode_splitfuse.xlsx")
                else:
                    excel_file = os.path.join(folder_path, "result_llama3.1-70b_decode.xlsx")
            
            write_to_excel(excel_file, txt_data)
            print(f"{file_name} writing to excel {excel_file}")
            read_file(full_path)

def main(txt_file, excel_file):
    """主函数，读取 txt 文件并将数据写入 Excel"""
    txt_data = read_txt_file(txt_file)
    print(txt_data)
    if txt_data:
        write_to_excel(excel_file, txt_data)

# 示例使用

# folder_path = "result"
# process_folder(folder_path)


# txt_file = 'result/qwen2.5-7b_pso_result_6_2_decode_prefix_in256_out256.txt'  # 假设你的文本文件名为 data.txt
# excel_file = 'result/result.xlsx'  # 生成的 Excel 文件名
# read_file(excel_file)
# main(txt_file, excel_file)
# read_file(excel_file)